from collections.abc import Iterator
from typing import Any, TypeVar, TypedDict

from zkl_aiutils_datasets.basics import Dataset, DatasetIterator, pause, resume, skip

AnySample = TypeVar('AnySample')


class SkippedIteratorState(TypedDict):
    iterator_state: Any
    skipped: bool


class SkippedIterator(DatasetIterator[AnySample]):
    def __init__(self, iterator: Iterator[AnySample], skip_samples_n: int, skipped: bool = False):
        self._skipped = skipped
        self._iterator = iter(iterator)
        self._skip_samples_n = skip_samples_n

    def __next__(self) -> AnySample:
        if not self._skipped:
            skip(self._iterator, self._skip_samples_n)
            self._skipped = True
        return next(self._iterator)

    def __skip__(self, samples_n: int):
        if not self._skipped:
            skip(self._iterator, self._skip_samples_n)
            self._skipped = True
        skip(self._iterator, samples_n)

    def __pause__(self) -> Any:
        return SkippedIteratorState(
            iterator_state=pause(self._iterator),
            skipped=self._skipped)


class SkippedDataset(Dataset[AnySample]):
    def __init__(self, dataset: Dataset[AnySample], skip_samples_n: int):
        self._dataset = dataset
        self._skip_samples_n = skip_samples_n

    @property
    def dataset(self) -> Dataset[AnySample]:
        return self._dataset

    @property
    def skip_samples_n(self) -> int:
        return self._skip_samples_n

    def __len__(self) -> int:
        return max(len(self._dataset) - self._skip_samples_n, 0)

    def __getitem__(self, index: int, /) -> AnySample:
        return self._dataset[index + self._skip_samples_n]

    def __iter__(self) -> SkippedIterator[AnySample]:
        return SkippedIterator(
            iterator=iter(self._dataset),
            skip_samples_n=self._skip_samples_n)

    def __resume__(self, state: Any | None = None, /) -> SkippedIterator[AnySample]:
        return SkippedIterator(
            iterator=resume(self._dataset, state['iterator_state']),
            skip_samples_n=self._skip_samples_n,
            skipped=state['skipped'])
